from typing import Dict, Optional, List
import re
import pymysql
from pymysql.cursors import DictCursor
from config import target_database


def results_to_markdown(results: List[Dict]) -> str:
    """
    将查询结果转换为Markdown表格
    :param results: 查询结果列表，每个元素为字典
    :return: Markdown表格字符串
    """
    if not results:
        return "无查询结果"

    # 获取表头（所有字典的键）
    headers = results[0].keys()

    # 构建表头行
    markdown = "| " + " | ".join(headers) + " |\n"

    # 构建分隔线
    markdown += "| " + " | ".join(["---"] * len(headers)) + " |\n"

    # 构建数据行
    for row in results:
        values = [str(row[header]) for header in headers]
        markdown += "| " + " | ".join(values) + " |\n"

    return markdown


class MySQLQueryExecutor:
    def __init__(self, host, user, password, database, port, charset='utf8mb4'):
        """初始化数据库连接参数"""
        self.host = host
        self.user = user
        self.password = password
        self.database = database
        self.port = port
        self.charset = charset
        self.connection = None
        self.cursor = None

    def connect(self):
        """建立数据库连接"""
        try:
            self.connection = pymysql.connect(
                host=self.host,
                user=self.user,
                password=self.password,
                database=self.database,
                port=self.port,
                charset=self.charset,
                cursorclass=DictCursor
            )
            self.cursor = self.connection.cursor()
            return True
        except pymysql.MySQLError as e:
            print(f"数据库连接失败: {e}")
            return False

    def execute_query(self, sql, params=None):
        """
        执行查询语句
        :param sql: 查询SQL语句.
        :param params: SQL参数，元组或列表形式，用于防止SQL注入.
        :return: 查询结果列表，每个元素为字典；失败返回空列表.
        """
        if not self.connection:
            print("请先连接数据库")
            return []

        try:
            # 执行查询，支持参数化查询
            self.cursor.execute(sql, params or ())

            # 获取查询结果
            results = self.cursor.fetchall()
            print(f"查询成功，返回 {len(results)} 条记录")
            return results

        except pymysql.MySQLError as e:
            print(f"查询执行失败: {e}")
            return []

    def get_tables(self, schema: Optional[str] = None) -> List[str]:
        """
        获取指定库下的所有表名
        :param schema: 库名；默认使用初始化时的 database
        :return: 表名列表
        """
        target_schema = schema or self.database
        sql = (
            """
            SELECT table_name
            FROM information_schema.tables
            WHERE table_schema = %s
            ORDER BY table_name
            """
        )
        rows = self.execute_query(sql, (target_schema,))
        return [row["TABLE_NAME"] for row in rows]

    def get_table_columns(self, table_name: str, schema: Optional[str] = None) -> List[Dict]:
        """
        获取指定表的字段信息
        :param table_name: 表名
        :param schema: 库名；默认使用初始化时的 database
        :return: 字段信息字典列表
        """
        target_schema = schema or self.database
        sql = (
            """
            SELECT
                COLUMN_NAME   AS column_name,
                COLUMN_TYPE   AS column_type,
                DATA_TYPE     AS data_type,
                CHARACTER_MAXIMUM_LENGTH AS char_max_length,
                IS_NULLABLE   AS is_nullable,
                COLUMN_DEFAULT AS column_default,
                COLUMN_KEY    AS column_key,
                EXTRA         AS extra,
                COLUMN_COMMENT AS column_comment,
                ORDINAL_POSITION AS ordinal_position
            FROM information_schema.columns
            WHERE table_schema = %s AND table_name = %s
            ORDER BY ORDINAL_POSITION
            """
        )
        return self.execute_query(sql, (target_schema, table_name))

    def _quote_identifier(self, name: str) -> str:
        """
        安全包裹标识符（库名/表名/列名）。仅允许字母、数字、下划线，避免注入。
        """
        if not re.match(r"^[A-Za-z0-9_]+$", name):
            raise ValueError(f"不安全的标识符: {name}")
        return f"`{name}`"

    def show_create_table(self, table_name: str, schema: Optional[str] = None) -> str:
        """
        获取建表语句（DDL）
        :param table_name: 表名
        :param schema: 库名；默认使用初始化时的 database
        :return: 建表语句字符串
        """
        if not self.connection:
            print("请先连接数据库")
            return ""

        target_schema = schema or self.database
        q_schema = self._quote_identifier(target_schema)
        q_table = self._quote_identifier(table_name)
        ddl_sql = f"SHOW CREATE TABLE {q_schema}.{q_table}"

        try:
            self.cursor.execute(ddl_sql)
            row = self.cursor.fetchone()
            # 不同 MySQL 版本键名可能不同，通常为 'Create Table'
            if not row:
                return ""
            for key in ("Create Table", "Create Table "):
                if key in row:
                    return row[key]
            # 兜底：返回第一列值
            return list(row.values())[1] if len(row.values()) > 1 else list(row.values())[0]
        except pymysql.MySQLError as e:
            print(f"SHOW CREATE TABLE 执行失败: {e}")
            return ""

    def get_database_structure(self, schema: Optional[str] = None) -> Dict[str, List[Dict]]:
        """
        获取整个库的结构：{ 表名: [字段信息...] }
        :param schema: 库名；默认使用初始化时的 database
        :return: 映射字典
        """
        structure: Dict[str, List[Dict]] = {}
        tables = self.get_tables(schema)
        for table in tables:
            structure[table] = self.get_table_columns(table, schema)
        return structure

    def structure_to_markdown(self, schema: Optional[str] = None) -> str:
        """
        将库的结构转为 Markdown（逐表输出字段信息）
        """
        structure = self.get_database_structure(schema)
        if not structure:
            return "无结构信息"

        parts: List[str] = []
        for table_name, columns in structure.items():
            parts.append(f"### 表：{table_name}")
            if columns:
                parts.append(results_to_markdown(columns))
            else:
                parts.append("无字段信息")
        return "\n\n".join(parts)

    def close(self):
        """关闭数据库连接"""
        if self.cursor:
            self.cursor.close()
        if self.connection:
            self.connection.close()
            print("数据库连接已关闭")


# if __name__ == "__main__":
#     # 初始化查询器
#     query_executor = MySQLQueryExecutor(
#         host=target_database["host"],
#         user=target_database["user"],
#         port=target_database["port"],
#         password=target_database["password"],
#         database=target_database["database"]
#     )
#
#     # 连接数据库
#     if query_executor.connect():
        # 示例1：简单查询并转换为Markdown
        # print("=== 所有用户 (Markdown格式) ===")
        # results = query_executor.execute_query("SELECT * FROM models LIMIT 10")
        # markdown_table = results_to_markdown(results)
        # print(markdown_table)
        #
        # # 关闭连接
        # query_executor.close()

        # if query_executor.connect():
        #     # 1) 列出所有表
        #     # print(query_executor.get_tables())
        #     #
        #     # # 4) 获取整库结构（字典）
        #     # structure = query_executor.get_database_structure()
        #     # print(structure)
        #
        #     # 5) 整库结构转 Markdown
        #     print(query_executor.structure_to_markdown())
        #
        #     query_executor.close()
